#!/usr/bin/env python3
# encoding: utf-8
# @data:2023/07/13
# @author:aiden
# 绕障碍
import rospy
import math
import signal
from ainex_sdk import pid, misc, common
from ainex_kinematics.gait_manager import GaitManager
from ainex_example.color_common import Common
from ainex_example.pid_track import PIDTrack
from ainex_example.visual_patrol_match import VisualPatrolMatch
from ainex_example.approach_object import ApproachObject
from ainex_interfaces.srv import SetString
from ainex_interfaces.msg import ObjectsInfo, ColorDetect, ROI

class CombinationNode(Common):
    # 按顺序检测三个roi，如果检测到黑线立刻跳出
    # y_min, y_max, x_min, x_max分别表示占图像的比例, 即实际大小为y_min*height
    line_roi = [(5 / 12, 6 / 12, 0, 1),
                (6 / 12, 7 / 12, 0, 1),
                (7 / 12, 8 / 12, 0, 1)
                ]

    column_roi = [0, 4 / 5, 0, 1]
    hurdles_roi = [1 / 5, 1, 0, 1]
    ball_roi = [0, 1, 0, 1] 
    # 所需动作的名称
    hurdles_action_name = 'hurdles'
    climb_stairs_action_name = 'climb_stairs_competition_2'
    descend_stairs_action_name = 'descend_stairs'
    back = 'go_back'
    go = 'forward_step'

    # 图像处理时缩放到这个分辨率， 不建议修改
    image_process_size = [160, 120]

    def __init__(self, name):
        rospy.init_node(name)
        self.name = name
        self.rl_dis = None
        self.current_color_1 = None
        self.current_color_2 = None
        self.current_color_3 = None
        self.first_location = None
        self.color_num = 0
        self.gait_manager = GaitManager()
        self.go_gait_param = self.gait_manager.get_gait_param()
        self.go_gait_param['body_height'] = 0.025
        self.go_gait_param['step_height'] = 0.02
        self.go_gait_param['hip_pitch_offset'] = 15
        self.go_gait_param['z_swap_amplitude'] = 0.006
        self.go_gait_param['init_roll_offset'] = 0
        self.go_dsp = [250, 0.2, 0.022]
        self.go_arm_swap = 30

        self.count = 0
        self.running = True
        self.objects_info = []

        self.current_state = "look"

        self.state = {'look':   [[500, 500], ['red', self.column_roi, self.image_process_size, self.set_column_color], False],
                      'walk':   [[500, 400], ['white', self.line_roi, self.image_process_size, self.set_visual_patrol_color], False],
                      'red':   [[500, 260], ['red', self.column_roi, self.image_process_size, self.set_column_color], False],
                      'yellow':   [[500, 260], ['yellow', self.column_roi, self.image_process_size, self.set_column_color], False],
                      'blue':   [[500, 260], ['blue', self.column_roi, self.image_process_size, self.set_column_color], False],
                      'ball':   [[500, 450], ['blue', self.ball_roi, self.image_process_size, self.set_ball_color], False],
                      }
        self.head_pan_init = 500  # 左右舵机的初始值
        self.head_tilt_init = 400 # 上下舵机的初始值
        super().__init__(name, self.head_pan_init, self.head_tilt_init)

        self.calib_config = common.get_yaml_data('/home/ubuntu/ros_ws/src/ainex_example/config/calib.yaml')

        self.approach_object = ApproachObject(self.gait_manager)
        self.visual_patrol_match = VisualPatrolMatch(self.gait_manager)
        self.visual_patrol_match.update_go_gait(dsp=self.go_dsp,walking_param=self.go_gait_param)
        
        # 头部的pid追踪
        self.head_pan_range = [125, 875]  # 左右转动限制在这个范围， 125为右
        self.head_tilt_range = [260, 500]  # 上下限制在这个范围， 250为下
        self.pid_rl = pid.PID(0.1, 0.0, 0.001)
        self.pid_ud = pid.PID(0.1, 0.0, 0.001)
        self.head_pan_init = 500  # 左右舵机的初始值
        self.head_tilt_init = 400 # 上下舵机的初始值
        self.rl_track = PIDTrack(self.pid_rl, self.head_pan_range, self.head_pan_init)
        self.ud_track = PIDTrack(self.pid_ud, self.head_tilt_range, self.head_tilt_init)

  # 躯体的追踪参数
        self.yaw_stop = 40  # 躯体停止转动时头部左右舵机脉宽值和中位500的差值

        signal.signal(signal.SIGINT, self.shutdown)
        
        # 订阅颜色识别结果
        rospy.Subscriber('/object/pixel_coords', ObjectsInfo, self.get_color_callback)
        rospy.Service('~set_color', SetString, self.set_color_srv_callback)  # 设置颜色
        self.motion_manager.run_action('walk_ready')

        if rospy.get_param('~start', True):
            # 通知颜色识别准备，此时只显示摄像头原画
            self.enter_func(None)
            self.start_srv_callback(None)
            common.loginfo('start ombination')

    def shutdown(self, signum, frame):
        self.running = False
        common.loginfo('%s shutdown' % self.name)

    def set_visual_patrol_color(self, color, roi, image_process_size):
        # 设置巡线颜色
        line_param = ColorDetect()
        line_param.color_name = color
        line_param.use_name = True
        line_param.detect_type = 'line'
        line_param.image_process_size = image_process_size
        line_param.line_roi.up.y_min = int(roi[0][0] * image_process_size[1])
        line_param.line_roi.up.y_max = int(roi[0][1] * image_process_size[1])
        line_param.line_roi.up.x_min = int(roi[0][2] * image_process_size[0])
        line_param.line_roi.up.x_max = int(roi[0][3] * image_process_size[0])

        line_param.line_roi.center.y_min = int(roi[1][0] * image_process_size[1])
        line_param.line_roi.center.y_max = int(roi[1][1] * image_process_size[1])
        line_param.line_roi.center.x_min = int(roi[1][2] * image_process_size[0])
        line_param.line_roi.center.x_max = int(roi[1][3] * image_process_size[0])

        line_param.line_roi.down.y_min = int(roi[2][0] * image_process_size[1])
        line_param.line_roi.down.y_max = int(roi[2][1] * image_process_size[1])
        line_param.line_roi.down.x_min = int(roi[2][2] * image_process_size[0])
        line_param.line_roi.down.x_max = int(roi[2][3] * image_process_size[0])

        line_param.min_area = 1
        line_param.max_area = image_process_size[0] * image_process_size[1]
        
        return line_param

    def set_column_color(self, color, roi, image_process_size):
        # 设置柱子标志颜色
        stairs_param = ColorDetect()
        stairs_param.color_name = color
        stairs_param.detect_type = 'side'
        stairs_param.use_name = True
        stairs_param.image_process_size = image_process_size
        stairs_param.roi.y_min = int(roi[0] * image_process_size[1])
        stairs_param.roi.y_max = int(roi[1] * image_process_size[1])
        stairs_param.roi.x_min = int(roi[2] * image_process_size[0])
        stairs_param.roi.x_max = int(roi[3] * image_process_size[0])
        stairs_param.min_area = 10*20
        stairs_param.max_area = image_process_size[0]*image_process_size[1]
        
        return stairs_param

    def set_ball_color(self, color, roi, image_process_size):
        # 设置追踪颜色
        ball_param = ColorDetect()
        ball_param.color_name = color
        ball_param.detect_type = 'circle'
        ball_param.use_name = True
        ball_param.image_process_size = self.image_process_size
        ball_param.roi.y_min = int(roi[0] * image_process_size[1])
        ball_param.roi.y_max = int(roi[1] * image_process_size[1])
        ball_param.roi.x_min = int(roi[2] * image_process_size[0])
        ball_param.roi.x_max = int(roi[3] * image_process_size[0])
        ball_param.min_area = 20
        ball_param.max_area = self.image_process_size[0]*self.image_process_size[1]

        return ball_param

    def set_color_srv_callback(self, msg):
        # 设置颜色
        stairs_param = self.set_column_color(self.state['climb_stairs'][1][0])
        line_param = self.set_visual_patrol_color(self.state['visual_patrol'][1][0])
        hurdles_param = self.set_column_color(self.state['hurdles'][1][0])

        self.detect_pub.publish([line_param, stairs_param, hurdles_param])
        common.loginfo('%s set_color' % self.name)
        
        return [True, 'set_color']

    def get_color_callback(self, msg):
        # 获取颜色识别结果
        red = False
        yellow = False
        blue = False
        self.objects_info = msg.data
        for color in self.objects_info:
            if self.current_state == 'look':
                if color.type == 'side':
                    self.current_color_1 = color.label  
                if self.current_color_1 is not None:
                    self.current_state = 'walk'
            if self.current_state == 'walk':     
                if color.type == 'side':   
                    if color.label != self.current_color_1 and self.current_color_2 is None:
                        self.current_color_2 = color.label
                    if self.current_color_2 is not None and self.current_color_3 is None:
                        if self.current_color_1 == 'red' and self.current_color_2 == 'yellow':
                            self.current_color_3 = 'blue'
                        if self.current_color_1 == 'red' and self.current_color_2 == 'blue':
                            self.current_color_3 = 'yellow'
                        if self.current_color_1 == 'yellow' and self.current_color_2 == 'blue':
                            self.current_color_3 = 'red'
                        if self.current_color_1 == 'yellow' and self.current_color_2 == 'red':
                            self.current_color_3 = 'blue'
                        if self.current_color_1 == 'blue' and self.current_color_2 == 'yellow':
                            self.current_color_3 = 'red'
                        if self.current_color_1 == 'blue' and self.current_color_2 == 'red':
                            self.current_color_3 = 'yellow'

            if self.current_state == 'ball':
                if self.objects_info is not None:
                    if color.type == 'circle':
                        color.x = color.x - self.calib_config['center_x_offset']
                        self.rl_dis, self.ud_dis = self.head_track_process(color)
                        if self.rl_dis is not None:
                            self.body_track_process(self.rl_dis, self.ud_dis, color)
                            self.rl_dis = None
                            self.count_miss = 0
                            self.start_find_ball = False

                    
    def state_init(self, current_state):
    # 不同阶段的初始化
        if self.state[current_state][2] == False:
            self.state[current_state][2] = True
            self.init_action(self.state[current_state][0][0], self.state[current_state][0][1])  # 头部姿态
            param1 = self.state[current_state][1][3](self.state[current_state][1][0], self.state[current_state][1][1], self.state[current_state][1][2])
            #param1 = self.state['walk'][1][3](self.state['walk'][1][0], self.state['walk'][1][1], self.state['walk'][1][2])
            self.detect_pub.publish([param1])  # 颜色检测设置
            common.loginfo(current_state + ' init')

    def state_init_look(self, current_state):
        # 不同阶段的初始化
        if self.state[current_state][2] == False:
            self.state[current_state][2] = True
            self.init_action(self.state[current_state][0][0], self.state[current_state][0][1])  # 头部姿态
            param2 = self.state['red'][1][3](self.state['red'][1][0], self.state['red'][1][1], self.state['red'][1][2])
            param3 = self.state['yellow'][1][3](self.state['yellow'][1][0], self.state['yellow'][1][1], self.state['yellow'][1][2])
            param4 = self.state['blue'][1][3](self.state['blue'][1][0], self.state['blue'][1][1], self.state['blue'][1][2])
            self.detect_pub.publish([param2, param3, param4])  # 颜色检测设置
            common.loginfo(current_state + ' init')

    def state_init_walk(self, current_state):
        # 不同阶段的初始化
        if self.state[current_state][2] == False:
            self.state[current_state][2] = True
            self.init_action(self.state[current_state][0][0], self.state[current_state][0][1])  # 头部姿态
            param1 = self.state[current_state][1][3](self.state[current_state][1][0], self.state[current_state][1][1], self.state[current_state][1][2])
            param2 = self.state['red'][1][3](self.state['red'][1][0], self.state['red'][1][1], self.state['red'][1][2])
            param3 = self.state['yellow'][1][3](self.state['yellow'][1][0], self.state['yellow'][1][1], self.state['yellow'][1][2])
            param4 = self.state['blue'][1][3](self.state['blue'][1][0], self.state['blue'][1][1], self.state['blue'][1][2])
            self.detect_pub.publish([param1, param2, param3, param4])  # 颜色检测设置
            common.loginfo(current_state + ' init')

    # 头部PID跟踪
    def head_track_process(self, object_info):
        # 头部追踪
        if abs(object_info.x - object_info.width/2) < 10:
            object_info.x = object_info.width/2
        if abs(object_info.y - object_info.height/2) < 10:
            object_info.y = object_info.height/2
        rl_dis = self.rl_track.track(object_info.x, object_info.width/2)
        ud_dis = self.ud_track.track(object_info.y, object_info.height/2)
        self.motion_manager.set_servos_position(20, [[23, int(rl_dis)], [24, int(ud_dis)]])
        return rl_dis, ud_dis
    
    # 躯体追踪
    def body_track_process(self, rl_dis, ud_dis, ball_data):                                                
        # 左右根据头部左右舵机位置值进行调整
        if ball_data is not None:
            if abs(rl_dis - 500) < 20:
                yaw_output =0
            else:
                yaw_output = math.copysign(5, rl_dis - 500)
            self.gait_manager.set_step(self.go_dsp, 0.02, 0, yaw_output, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)

    # 主循环函数
    def run(self):
        start = True
        end_color = False
        end = False
        while self.running:
            if start:
                # 获取识别结果
                line_data = None
                column_data = None
                # stairs_data = None
                for object_info in self.objects_info:
                    if object_info.type == 'line':
                        line_data = object_info
                    if object_info.type == 'side':
                        column_data = object_info

                # 如果有识别结果,巡线控制
                if self.current_state == 'walk':
                    print(self.current_color_1, self.current_color_2, self.current_color_3)
                    if line_data is not None:
                        self.visual_patrol_match.process(line_data.x, line_data.width)
                        self.line_data = None
                    elif column_data is not None:
                        if column_data.label == self.current_color_1 or column_data.label == self.current_color_3:
                            self.gait_manager.set_step(self.go_dsp, 0.02, 0, 8, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                            if column_data.label == self.current_color_3:
                                end_color = True
                                self.count = 0
                        elif column_data.label == self.current_color_2:
                            self.gait_manager.set_step(self.go_dsp, 0.02, 0, -8, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)

                    elif line_data is None and column_data is None:
                        self.gait_manager.set_step(self.go_dsp, 0.02, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                    if end_color and column_data is None:
                        self.count += 1
                        print(self.count)
                        if self.count > 1600:
                            self.current_state = 'ball'
                            self.color_num = 0
                            self.count = 0
     
                rospy.sleep(0.01)
                if self.current_state == 'look':
                    self.state_init_look(self.current_state)
                if self.current_state == 'walk':
                    self.state_init_walk(self.current_state)
                if self.current_state == 'ball':
                    self.state_init(self.current_state)
            else:
                rospy.sleep(0.01)
            
        # 退出前动作
        self.init_action(self.head_pan_init, self.head_tilt_init)
        self.stop_srv_callback(None)
        rospy.signal_shutdown('shutdown')

if __name__ == "__main__":
    CombinationNode('obstacle').run()
